#include <iostream>
#include "trackHandle.h"
#include "TFile.h"
#include "TChain.h"

using namespace std;
using namespace cv;

const int trackNX = 72;
const int trackNY = 72;
const double pixelSize = 0.080;
const double xAxisLength = trackNX*pixelSize;
const double yAxisLength = trackNY*pixelSize;
const double Ecal = 5.9;
const double EcalADC = 2584;
const double ErrorCnt = 0.267;

Double_t cacu_ellip(short data[trackNX][trackNY]){
    double u11=0;
    double u20=0;
    double u02=0;
    double M2max=0;
    double M2min=0;
    double eli=1;
    double mBx=0;
    double mBy=0;
    double qtotl = 0;
    for(int i =0; i<trackNX;i++){
        for(int j =0; j<trackNY;j++){
            qtotl+=double(data[i][j]);
            mBx+=double(data[i][j])*i;
            mBy+=double(data[i][j])*j;
        }
    }
    mBx /= qtotl;
    mBy /= qtotl;
    for(int i =0; i<trackNX;i++){
        for(int j =0; j<trackNY;j++){
            double xx=i;
            double yy = j;
            double qq = double(data[i][j]);
            u11+=(xx-mBx)*(yy-mBy)*qq;
            u20+=pow(xx-mBx,2)*qq;
            u02+=pow(yy-mBy,2)*qq;
        }
    }
    M2max=0.5*(u20+u02+sqrt(pow(u20-u02,2)+4*pow(u11,2)));
    M2min=0.5*(u20+u02-sqrt(pow(u20-u02,2)+4*pow(u11,2)));
    eli=sqrt(1-M2min/M2max);
    return eli;
}

Double_t cacu_Oirc(short data[trackNX][trackNY]){
    double xx=0;
    double yy=0;
    double qq=0;
    double m10=0;
    double m01=0;
    double m00=0;
    double u00=0;
    double u20=0;
    double u02=0;
    double cir=0;
    double adc=0;
    double pi=M_PI;
    for(int i=0;i<trackNX;i++){
        for(int j =0;j<trackNY;j++){
            xx = i;
            yy = j;
            qq = double(data[i][j]);
            if (qq>0)
            {
                m10+=xx;
                m01+=yy;
                m00+=1;
            }
        }
    }
    for(int i=0;i<trackNX;i++){
        for(int j =0;j<trackNY;j++){
            xx = i;
            yy = j;
            qq = double(data[i][j]);
            if (qq>0)
            {
               u20+=pow((xx-(m10/m00)),2);
                u02+=pow((yy-(m01/m00)),2);
                u00+=1;
                adc+=qq;
            }
        }
    }
    cir=pow(u00,2)/((u20+u02)*2*pi);
    return cir;
}

Double_t FindMin(short data[trackNX][trackNY]){
    double MinPixel = 9999;
    double qq;
    for(int i=0;i<trackNX;i++){
        for(int j =0;j<trackNY;j++){
            qq = double(data[i][j]);
            if (qq>2)
            {
                if(qq<MinPixel){
                    MinPixel = qq;
                }
            }
        }
    }
    return MinPixel>35? 35:MinPixel;
}

Double_t EdpDense(short data[trackNX][trackNY]){
    double PixelNum = 0.;
    double PixelSum = 0.;
    for(int i=0; i<trackNX; i++){
        for(int j =0; j<trackNY;j++){
            if(data[i][j]>1){
                PixelNum+=1.;
                PixelSum+=double(data[i][j]);
            }
        }
    }
    return PixelSum/PixelNum;
}

int FindStartEdgeX(short data[trackNX][trackNY]){
    int edgx = 0;
    for(int ix = 0; ix<trackNX; ix++){
        for(int iy = 0; iy<trackNY; iy++){
            if(data[ix][iy]>20){
                edgx = ix;
                return edgx;
            }
        }
    }
    return edgx;
}

int FindEndEdgeX(short data[trackNX][trackNY]){
    int edgx = trackNX-1;
    for(int ix = edgx; ix>=0; ix--){
        for(int iy = edgx; iy>=0; iy--){
            if(data[ix][iy]>20){
                edgx = ix;
                return edgx;
            }
        }
    }
    return edgx;
}

int FindEdgeLen(short data[trackNX][trackNY], int edgestart){
    int len =0;
    for(int iy = 0; iy<trackNY; iy++){
        if(data[edgestart][iy]>0){
            len++;
        }
    }
    return len;
}

int FindMaxlenX(short data[trackNX][trackNY], int edgesatrt, int edgeend){
    int maxlen = 0;
    for(int ix = edgesatrt; ix<=edgeend; ix++){
        int tmplen = 0;
        for(int iy = 0 ; iy<trackNY; iy++){
            if(data[ix][iy]>0){tmplen+=1;}
        }
        if(maxlen<tmplen){
            maxlen=tmplen;
        }
    }
    return maxlen;
}

double StartXedgeGradient(short data[trackNX][trackNY], int edgesatrt){
    int SumDiff;
    int Hitnum = 0;
    if(edgesatrt==0){
        for(int iy = 0; iy<trackNY; iy++){
            SumDiff += data[edgesatrt][iy];
            if(data[edgesatrt][iy]>0) Hitnum+=1;
        }
    }
    else{
        for(int iy = 0; iy<trackNY; iy++){
            SumDiff += (data[edgesatrt][iy]-data[edgesatrt-1][iy]);
            if(data[edgesatrt][iy]>0) Hitnum+=1;
        }
    }
    return double(SumDiff)/double(Hitnum);

}

double EndXedgeGradient(short data[trackNX][trackNY], int edgeEnd){
    int SumDiff;
    int Hitnum = 0;
    if(edgeEnd==71){
        for(int iy = 0; iy<trackNY; iy++){
            SumDiff += data[edgeEnd][iy];
            if(data[edgeEnd][iy]>0) Hitnum+=1;
        }
    }
    else{
        for(int iy = 0; iy<trackNY; iy++){
            SumDiff += (data[edgeEnd][iy]-data[edgeEnd+1][iy]);
            if(data[edgeEnd][iy]>0) Hitnum+=1;
        }
    }
    return double(SumDiff)/double(Hitnum);

}

TrackHandle::TrackHandle(int Threshold) {
    ADC_Threshold = Threshold;
}

TrackHandle::~TrackHandle() {
    // vector<myTrack>().swap(Vtrack);
}

Mat TrackHandle::S16C1_to_U8C1(Mat S16C1_Array)
{
    double min, max;
    Mat U8C1_Array;
    minMaxIdx(S16C1_Array, &min, &max);
    S16C1_Array.convertTo(U8C1_Array, CV_8UC1, 255/(max-min), -min);

    return U8C1_Array;
}

Mat TrackHandle::Array2Mat(TH2F* hist, uint32_t w, uint32_t h) {
    Mat track_img(h, w, CV_16SC1, cv::Scalar(0));
    
    for (int i = 0; i < h; ++i) {
        for (int j = 0; j < w; ++j) {
            track_img.at<short>(i, j) =  hist->GetBinContent(i+1, j+1);
        }
    }

    track_img.at<short>(0, 0) = 0;
    track_img.at<short>(0, 1) = 0;
    track_img.at<short>(0, 2) = 0;
    track_img.at<short>(0, 3) = 0;

    return track_img;
}

int TrackHandle::Length(Mat frame_center) {
    Mat frame_center_blur;
    medianBlur(frame_center, frame_center_blur, 3);
    Mat frame_thinning;
    ximgproc::thinning(frame_center, frame_thinning);
    int frame_len = sum(frame_thinning)[0]/255;
    
    return frame_len;
}

Mat TrackHandle::RoiCenter(Mat frame, bool isSource) {
    Mat frame_blur;
    GaussianBlur(frame, frame_blur, Size(3, 3), 0);
    Mat frame_binary;
    threshold(frame_blur, frame_binary, 0, 255, cv::THRESH_BINARY);
    Mat frame_close;
    Mat kernel = getStructuringElement(MORPH_RECT,Size(2,2));
    morphologyEx(frame_binary, frame_close, MORPH_CLOSE, kernel);
    vector<vector<cv::Point>> contours;
    vector<cv::Point> max_contours;
    double area = 0.0;
    findContours(frame_close, contours, cv::RETR_EXTERNAL, cv::CHAIN_APPROX_NONE);
    for(auto &iter:contours) {
        double area_temp = contourArea(iter);
        if(area_temp >= area) {
            area = area_temp;
            max_contours = iter;
        }
    }
    if(area == 0.0) {
        contour_size = 0;
        if(isSource == false) {
            return frame_binary;
        } else {
            return frame;
        }
    }
    
    contour_size = max_contours.size();
    Rect border = boundingRect(max_contours);
    Mat ROI;
    if(isSource == false) {
        Mat track_fill(frame.size(), CV_8UC1, Scalar::all(0));
        std::vector<std::vector<cv::Point>> temp_contours;
        temp_contours.push_back(max_contours);
        drawContours(track_fill, temp_contours, -1, 255, -1);
        ROI = track_fill(border);
    } else {
        ROI = frame_blur(border);
    }
    
    int roi_x, roi_y, roi_w, roi_h;
    roi_x = border.x;
    roi_y = border.y;
    roi_w = border.width;
    roi_h = border.height;

    track_w = roi_w;
    track_h = roi_h;

    int center_y = (trackNX-roi_h)/2;
    int center_x = (trackNY-roi_w)/2;

    Mat track_center(frame.size(), CV_8UC1, Scalar::all(0));
    for(int y = center_y; y < center_y+roi_h; y++) {
        for(int x = center_x; x < center_x+roi_w; x++) {
            track_center.at<uchar>(y, x) = ROI.at<uchar>(y-center_y, x-center_x);
        }
    }

    return track_center;
}

Mat TrackHandle::Hit_Frame(Mat frame) {
    Mat hit_frame;
    threshold(frame, hit_frame, ADC_Threshold, 32767, cv::THRESH_TOZERO);
    return hit_frame;
}

int TrackHandle::Max_Value(Mat frame) {
    double min, max;
    minMaxIdx(frame, &min, &max);

    return max;
}

coordinate TrackHandle::Max_Coor(Mat frame) {
    coordinate Peak;
    double minVal;
    double maxVal;
    Point minLoc;
    Point maxLoc;

    minMaxLoc(frame, &minVal, &maxVal, &minLoc, &maxLoc, Mat());
    Peak.x = maxLoc.x;
    Peak.y = maxLoc.y;

    return Peak;
}

int TrackHandle::Sum_Value(Mat frame) {
    return sum(frame)[0];
}

int TrackHandle::Count_NonZero(Mat frame) {
    return countNonZero(frame);
}

double TrackHandle::Diagonal_Length(int x, int y) {
    return sqrt(pow(x, 2) + pow(y, 2));
}

double TrackHandle::Ratio_XY(int x, int y) {
    return x/y;
}

coordinate TrackHandle::Bcenter_Coor(Mat frame) {
    coordinate Bcenter;
    int cols = frame.cols;
    int rows = frame.rows;
    double sumXV = 0.0;
    double sumYV = 0.0;
    double sumVal = 0.0;
    for(int y = 0; y < rows; y++) {
        for(int x = 0; x < cols; x++) {
            short val = frame.at<short>(y,x);
            if(val != 0) {
                sumXV += val*x;
                sumYV += val*y;
            }
        }
    }
    sumVal = double(sum(frame)[0]);
    Bcenter.x = sumXV / sumVal;
    Bcenter.y = sumYV / sumVal;
    
    return Bcenter;
}

double TrackHandle::Dis_Bcenter_Max(coordinate max_coor, coordinate bcenter_coor) {
    return sqrt(pow(max_coor.x - bcenter_coor.x, 2) + pow(max_coor.y - bcenter_coor.y, 2));
}

int TrackHandle::Contour_Size() {
    return contour_size;
}

std::pair<double, double> TrackHandle::truthPixel(TVector3 &truthPosition)
{
    double xPos = (truthPosition[0] + xAxisLength / 2.) / pixelSize;
    double yPos = (truthPosition[1] + yAxisLength / 2.) / pixelSize;
    return std::pair<double, double>(xPos, yPos);
}

void TrackHandle::feature_extraction(TTree* tree, TTree* truthTree, TTree* outTree, TTree* simTree) {
    Long64_t numOfEntries = tree->GetEntries();
    TH2F* hist = nullptr;
    Int_t iEvent = 0;
    Int_t iTruthEvent = 0;
    Int_t currentTreeNum = -1;
    Int_t thisTreeTotalNum = 0;
    Int_t lastTreeTotalNum = 0;
    tree->SetBranchAddress("iEvent", &iEvent);
    tree->SetBranchAddress("TopMetalHist", &hist);
    tree->GetBranch("TopMetalHist")->SetAutoDelete(kTRUE);

    TVector3* TruthPosition = nullptr;
    truthTree->SetBranchAddress("TruthPosition", &TruthPosition);
    truthTree->GetBranch("TruthPosition")->SetAutoDelete(kTRUE);
    TVector3* PolPhotElectronPos = nullptr;
    truthTree->SetBranchAddress("iEvent", &iTruthEvent);
    truthTree->SetBranchAddress("PolPhotElectronPos", &PolPhotElectronPos);
    truthTree->GetBranch("PolPhotElectronPos")->SetAutoDelete(kTRUE);
    Double_t TruthEnergy;
    truthTree->SetBranchAddress("TruthEnergy", &TruthEnergy);

    SimEvent* simEvent = new SimEvent();
    if (simTree != nullptr)
    {
        simTree->SetBranchAddress("SimEvent", &simEvent);
        //simTree->GetBranch("SimEvent")->SetAutoDelete(true);
        TBranch* br_SimEvent = outTree->Branch("SimEvent", "SimEvent", &simEvent, 32000, 0);
        br_SimEvent->SetAutoDelete(true);
    }

    Short_t data[trackNX][trackNY];   outTree->Branch("data",&data,Form("data[%d][%d]/S", trackNX, trackNY));
    outTree->Branch("TruthEventNum",&iEvent,"TruthEventNum/I");
    Int_t numEvent = 0;     outTree->Branch("numEvent",&numEvent,"numEvent/I");
    Int_t numFrame = 0;     outTree->Branch("numFrame",&numFrame,"numFrame/I");
    Int_t peakADC = 0;      outTree->Branch("peakADC",&peakADC,"peakADC/I");
    Double_t peakX;         outTree->Branch("peakX",&peakX,"peakX/D");
    Double_t peakY;         outTree->Branch("peakY",&peakY,"peakY/D");
    Int_t sumADC;           outTree->Branch("sumADC",&sumADC,"sumADC/I");
    Int_t HitADC;           outTree->Branch("HitADC",&HitADC,"HitADC/I");
    Int_t sumHit;           outTree->Branch("sumHit",&sumHit,"sumHit/I");
    Double_t lenXY;         outTree->Branch("lenXY",&lenXY,"lenXY/D");
    Double_t ratioXY;       outTree->Branch("ratioXY",&ratioXY,"ratioXY/D");
    Double_t BcenterX;      outTree->Branch("BcenterX",&BcenterX,"BcenterX/D");
    Double_t BcenterY;      outTree->Branch("BcenterY",&BcenterY,"BcenterY/D");
    Double_t lenBtoP;       outTree->Branch("lenBtoP",&lenBtoP,"lenBtoP/D");
    Int_t length;           outTree->Branch("length",&length,"length/I");
    Double_t width;         outTree->Branch("width",&width,"width/D");
    Int_t outline;          outTree->Branch("outline",&outline,"outline/I");
    Int_t edgeHit;          outTree->Branch("edgeHit",&edgeHit,"edgeHit/I");
    Double_t ellip;         outTree->Branch("ellip",&ellip,"ellip/D");
    Double_t circ;          outTree->Branch("circ",&circ,"circle/D");
    Double_t minpixel;      outTree->Branch("minpixel",&minpixel,"minpixel/D");
    Double_t energydense;   outTree->Branch("energydense",&energydense,"energydense/D");
    Double_t EdgeLenRatio;  outTree->Branch("EdgeLenRatio",&EdgeLenRatio,"EdgeLenRatio/D");
    Double_t Gradient;      outTree->Branch("Gradient",&Gradient,"Gradient/D");
    Double_t dEdX;          outTree->Branch("dEdX",&dEdX,"dEdX/D");
    Double_t Edp;           outTree->Branch("Edp",&Edp,"Edp/D");
    Double_t EdpEr;         outTree->Branch("EdpEr",&EdpEr,"EdpEr/D");
    Double_t TruthPx;       outTree->Branch("TruthPx",&TruthPx,"TruthPx/D");
    Double_t TruthPy;       outTree->Branch("TruthPy",&TruthPy,"TruthPy/D");
    Double_t PhotPx;       outTree->Branch("PhotPx",&PhotPx,"PhotPx/D");
    Double_t PhotPy;       outTree->Branch("PhotPy",&PhotPy,"PhotPy/D");
    outTree->Branch("TruthEnergy",&TruthEnergy,"TruthEnergy/D");

    for(int i = 0; i < numOfEntries; i++) {
        tree->GetEntry(i);
        truthTree->GetEntry(iEvent);
        //std::cout << __LINE__ << "Truth Energy" << TruthEnergy << std::endl;
        if(simTree != nullptr)
        {
            simTree->GetEntry(iEvent);
            if (TruthEnergy != simEvent->GetTruthEnergy())
            {
                std::clog << "Warning! The truth energy from Sim and Digi file unmatched. Please check it!" << std::endl;
            }
            if (iEvent != iTruthEvent)
            {
                std::clog << "Warning! The truth entry unmatched. TopMetal: " << iEvent << "\tTruth: " << iTruthEvent << std::endl;
            }
        }
        if(i%1000 == 0)
        {
            std::cout << "Processing event " << i << " / " << numOfEntries << std::endl;
        }

        Mat frame = Array2Mat(hist, trackNX, trackNY);
        if(countNonZero(frame) < 1) continue;
        Mat frame_hit = Hit_Frame(frame);
        // Mat frame_u8 = S16C1_to_U8C1(frame);
        Mat frame_hit_u8 = S16C1_to_U8C1(frame_hit);
        Mat frame_center_u8 = RoiCenter(frame_hit_u8, false);

        for(int ix = 0; ix < trackNX; ix++)
        {
            for(int iy = 0; iy < trackNY; iy++)
            {
                data[ix][iy] = frame.at<short>(ix,iy);
            }
        }
        peakADC = Max_Value(frame_hit);
        coordinate Peak = Max_Coor(frame_hit);
        peakX =     Peak.x;
        peakY =     Peak.y;
        sumADC = Sum_Value(frame);
        HitADC = Sum_Value(frame_hit);
        sumHit = Count_NonZero(frame_hit);
        lenXY = Diagonal_Length(track_w, track_h);
        ratioXY = Ratio_XY(track_w, track_h);
        coordinate Bcenter = Bcenter_Coor(frame_hit);
        BcenterX =  Bcenter.x;
        BcenterY =  Bcenter.y;
        lenBtoP = Dis_Bcenter_Max(Peak, Bcenter);
        length = Length(frame_center_u8);
        width = double(sumHit) / double(length);
        outline = Contour_Size();
        std::pair<double, double> truthPos = truthPixel(*TruthPosition);
        TruthPx = truthPos.first;
        TruthPy = truthPos.second;
        std::pair<double, double> photPos = truthPixel(*PolPhotElectronPos);
        PhotPx = photPos.first;
        PhotPy = photPos.second;
        int edgeHitSum = 0;
        int numAdd = 0, numLoss = 0;
        for(int lin = 0; lin < trackNX; lin ++){
            for(int row = 0; row < trackNY; row ++){
                //track.data[lin][row] = frame_hit.at<short>(lin,row);
                if(frame_hit.at<short>(lin,row) != 0)
                {
                    if(lin == 0 || lin == 71 || row == 0 || row == 71)
                    {
                        edgeHitSum ++;
                    }
                }
                double subPixel = frame_hit.at<short>(lin,row) - preData[lin][row];
                if(subPixel > 20) numAdd++;
                if(subPixel < -20) numLoss++;
                preData[lin][row] = frame_hit.at<short>(lin,row);
            }
        }
        edgeHit = edgeHitSum;

        //------------for simulation----------------------
        numEvent ++;
        numFrame = 1;

        Edp = (double(sumADC)/(EcalADC))*Ecal;
        EdpEr = 0.267*sqrt(Edp);
        ellip = cacu_ellip(data);
        circ = cacu_Oirc(data);

        minpixel = FindMin(data);
        energydense = EdpDense(data);
        dEdX = 5.9*((double(sumADC))/double(length))/((EcalADC*(0.083)));
        int iedgestart = FindStartEdgeX(data);
        int iedgeend = FindEndEdgeX(data);
        int iedgestartlen = FindEdgeLen(data,iedgestart);
        int iedgeendtlen = FindEdgeLen(data,iedgeend);
        int MaxEdgeLen;
        if(iedgestartlen>iedgeendtlen){
            MaxEdgeLen = iedgestartlen;
        }
        else{
            MaxEdgeLen = iedgeendtlen;
        }
        int MaxLenx = FindMaxlenX(data,iedgestart,iedgeend);
        EdgeLenRatio = double(MaxEdgeLen)/double(MaxLenx);
        double startGrad = StartXedgeGradient(data,iedgestart);
        double endGrad = EndXedgeGradient(data,iedgeend);
        if(startGrad>endGrad){
            Gradient = startGrad;
        }
        else{
            Gradient = endGrad;
        }
        outTree->Fill();

    }
}


TTree* TrackHandle::calculate(TString filename, TString treeName, TString simFile) {
    TChain* tree = new TChain("TopMetal_0", "");
    TChain* truthTree = new TChain("TruthTree", "");

    std::cout << "Add file: " << filename.Data() << std::endl;
    tree->Add(filename.Data());
    truthTree->Add(filename.Data());

    TChain* simTree = nullptr;
    if (simFile != TString(""))
    {
        simTree = new TChain("Sim", "");
        simTree->Add(simFile.Data());
    }

    TTree* outTree = truthTree->CloneTree(0);
    outTree->SetName(treeName);

    feature_extraction(tree, truthTree, outTree, simTree);

    delete tree;
    delete truthTree;
    if(simTree != nullptr)
    {
        delete simTree;
    }

    return outTree;
}
